package main

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

// 设置应用配置后直接执行
var (
	baseUrl = "wss://aiui.xf-yun.com/v2/aiint/ws"
	// 应用配置
	appid     = ""
	apikey    = ""
	apisecret = ""
	// 场景
	scene = "main_box"

	// 用于设置发送文本请求还是音频请求
	// 请求数据类型：text/audio
	// 音频请求需要先设置filePath
	// 当前音频格式默认pcm，修改音频格式需要修改audioReq中的payload中音频相关参数
	dataType = "text"
	// 请求文本呢
	payLoadText = "你是谁"

	// 每帧音频数据大小，单位字节
	chunckSize int64 = 1024
	// 音频文件位置
	filePath = "D:/temp/audio/date.pcm"
)

func main() {
	// 设置握手超时
	d := websocket.Dialer{
		HandshakeTimeout: 5 * time.Second,
	}
	// 建立websocket连接
	handshakeUrl := assembleAuthUrl(baseUrl, apikey, apisecret)
	conn, resp, err := d.Dial(handshakeUrl, nil)
	if err != nil || resp.StatusCode != 101 {
		// 握手失败
		log.Panicf("handshake error：%v, %v", readResp(resp), err.Error())
	}

	done := make(chan struct{})
	go sendReq(conn, done)
	go receiveRes(conn, done)
	// 等待交互完成
	<-done
}

// 发送请求数据
func sendReq(conn *websocket.Conn, done chan struct{}) {
	var err error
	if dataType == "text" {
		req := genTextParam()
		log.Printf("send req: %v\n", marshalToString(req))
		err = conn.WriteJSON(req)
	}
	if dataType == "audio" {
		f, err := os.OpenFile(filePath, os.O_RDONLY, os.ModePerm)
		if err != nil {
			log.Printf("open file error: %v", err.Error())
			conn.Close()
			close(done)
			return
		}

		buf := make([]byte, chunckSize)
		fileInfo, _ := f.Stat()
		if err != nil {
			log.Printf("get file info error: %v", err.Error())
			conn.Close()
			close(done)
			return
		}
		fileSize := fileInfo.Size()
		frameMax := fileSize / chunckSize
		if fileSize%chunckSize != 0 {
			frameMax += 1
		}

		status := 0
		for i := int64(1); i <= frameMax; i++ {
			n, err := f.Read(buf)
			if err != nil && err != io.EOF {
				break
			}
			if frameMax == 1 {
				// 一帧传完
				status = 3
			} else if i == frameMax {
				// 尾帧
				status = 2
			} else if i != 1 {
				// 中间帧
				status = 1
			}

			req := genAudioParam(buf[:n], status)
			err = conn.WriteJSON(req)
			if err != nil {
				break
			}
		}
	}

	if err != nil {
		log.Printf("sendReq error: %v", err.Error())
		conn.Close()
		close(done)
	}
}

// 接收返回值
func receiveRes(conn *websocket.Conn, done chan struct{}) {
	defer close(done)
	for {
		_, d, err := conn.ReadMessage()
		if err != nil {
			log.Println("read message error: ", err.Error())
			return
		}
		// log.Println("response: ", string(d))

		aiuiRes := new(AIUIV2Res)
		err = json.Unmarshal(d, aiuiRes)
		if err != nil {
			fmt.Println("Unmarshal response error: ", err.Error())
			return
		}

		if aiuiRes.Header.Code != 0 {
			log.Println("response error: ", string(d))
			return
		}

		if aiuiRes.Payload.Iat != nil {
			dbytes, err := base64.StdEncoding.DecodeString(aiuiRes.Payload.Iat.Text)
			if err != nil {
				log.Println("Decode Text error: ", err.Error())
				return
			}
			log.Println("iat result: " + string(dbytes))
		}
		if aiuiRes.Payload.Nlp != nil {
			dbytes, err := base64.StdEncoding.DecodeString(aiuiRes.Payload.Nlp.Text)
			if err != nil {
				log.Println(aiuiRes.Payload.Nlp.Text)
				log.Println("Decode Text error: ", err.Error())
				return
			}
			log.Println("nlp result: " + string(dbytes))
		}
		if aiuiRes.Header.Status == 2 {
			// 最后一帧结果
			return
		}
	}
}

// 生成文本请求参数
func genTextParam() map[string]interface{} {
	return map[string]interface{}{
		"header": map[string]interface{}{
			"sn":     "4b37ccbb6679c34389176ed5bc920196",
			"appid":  appid,
			"stmid":  "text-1",
			"status": 3,
			"scene":  scene,
		},
		"parameter": map[string]interface{}{
			"nlp": map[string]interface{}{
				"nlp": map[string]interface{}{
					"compress": "raw",
					"format":   "json",
					"encoding": "utf8",
				},
				"sub_scene":   "cbm_v45",
				"new_session": true,
			},
		},
		"payload": map[string]interface{}{
			"text": map[string]interface{}{
				"compress": "raw",
				"format":   "plain",
				"text":     base64.StdEncoding.EncodeToString([]byte(payLoadText)),
				"encoding": "utf8",
				"status":   3,
			},
		},
	}
}

// 生成音频请求参数
func genAudioParam(data []byte, status int) map[string]interface{} {
	return map[string]interface{}{
		"header": map[string]interface{}{
			"sn":     "4b37ccbb6679c34389176ed5bc920196",
			"appid":  appid,
			"stmid":  "audio-1",
			"status": status,
			"scene":  scene,
		},
		"parameter": map[string]interface{}{
			"nlp": map[string]interface{}{
				"nlp": map[string]interface{}{
					"compress": "raw",
					"format":   "json",
					"encoding": "utf8",
				},
				"sub_scene":   "cbm_v45",
				"new_session": true,
			},
		},
		"payload": map[string]interface{}{
			"audio": map[string]interface{}{
				"encoding":    "raw",
				"sample_rate": 16000,
				"channels":    1,
				"bit_depth":   16,
				"frame_size":  len(data),
				"status":      status,
				"audio":       base64.StdEncoding.EncodeToString(data),
			},
		},
	}
}

// 创建鉴权url
func assembleAuthUrl(hosturl string, apikey, apisecret string) string {
	ul, err := url.Parse(hosturl)
	if err != nil {
		log.Panic("parse url error: ", err.Error())
	}
	// 签名时间 Sun, 03 Sep 2023 16:47:59 GMT
	date := time.Now().UTC().Format(time.RFC1123)
	date = strings.Replace(date, "UTC", "GMT", -1)
	// 参与签名的字段 host ,date, request-line
	signString := []string{"host: " + ul.Host, "date: " + date, "GET " + ul.Path + " HTTP/1.1"}
	// 拼接签名字符串
	sgin := strings.Join(signString, "\n")
	// fmt.Println(sgin)
	// 签名结果
	sha := HmacWithShaTobase64("hmac-sha256", sgin, apisecret)
	// 构建请求参数 此时不需要urlencoding
	authUrl := fmt.Sprintf("api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"", apikey,
		"hmac-sha256", "host date request-line", sha)
	// 将请求参数使用base64编码
	authorization := base64.StdEncoding.EncodeToString([]byte(authUrl))

	v := url.Values{}
	v.Add("host", ul.Host)
	v.Add("date", date)
	v.Add("authorization", authorization)
	// 将编码后的字符串url encode后添加到url后面
	return hosturl + "?" + v.Encode()
}

func HmacWithShaTobase64(algorithm, data, key string) string {
	mac := hmac.New(sha256.New, []byte(key))
	mac.Write([]byte(data))
	encodeData := mac.Sum(nil)
	return base64.StdEncoding.EncodeToString(encodeData)
}

func readResp(resp *http.Response) string {
	if resp == nil {
		return ""
	}
	b, err := io.ReadAll(resp.Body)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("code=%d,body=%s", resp.StatusCode, string(b))
}

func marshalToString(r interface{}) string {
	b, _ := json.Marshal(r)
	return string(b)
}

type AIUIV2Res struct {
	Header  AIUIV2ResHeader  `json:"header"`
	Payload AIUIV2ResPayload `json:"payload"`
}

type AIUIV2ResHeader struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
	Cid     string `json:"cid"`
	Sid     string `json:"sid"`
	Stmid   string `json:"stmid"`
	Status  int    `json:"status"`
}

type AIUIV2ResPayload struct {
	Iat *AIUIV2ResPayloadCommonData `json:"iat"`
	Nlp *AIUIV2ResPayloadCommonData `json:"nlp"`
}

type AIUIV2ResPayloadCommonData struct {
	Encoding string `json:"encoding"`
	Compress string `json:"compress"`
	Format   string `json:"format"`
	Status   int    `json:"status"`
	Text     string `json:"Text"`
	Seq      int    `json:"seq"`
}
